package com.firexxx.study.transformers;

import com.firexxx.study.annotation.TimeCost;
import javassist.*;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * @description: TODO
 * @throws:
 * @author: zhuangzhang2
 * @date: 2020/07/02 12:27
 */
public class CustomClassFileTransformer implements ClassFileTransformer {

    private String classPrefix = null;

    public CustomClassFileTransformer() {
    }

    public CustomClassFileTransformer(String classPrefix) {
        this.classPrefix = classPrefix;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
            if (className == null || className.length() == 0) {
                return null;
            }

            className = className.replace("/", ".");
            if (classPrefix != null && !className.startsWith(classPrefix)) {
                return null;
            }

            ClassPool pool = ClassPool.getDefault();
            CtClass ctClass = pool.get(className);

            boolean exist = false;

            for (CtMethod ctMethod : ctClass.getDeclaredMethods()) {
                if (ctMethod.getAnnotation(TimeCost.class) != null) {
                    String prompt = "[" + ctClass.getName() + "][" + ctMethod.getLongName() + "]";

                    System.out.println("发现需要转换的方法" + prompt);

                    enhance(ctClass, ctMethod);

                    exist = true;
                }
            }

            return exist ? ctClass.toBytecode() : null;
//            return null;
        } catch (NotFoundException e) {
//            e.printStackTrace();
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }

    private void enhance(CtClass ctClass, CtMethod ctMethod) throws CannotCompileException, IOException {
        String prefix = "\nlong startTime = System.currentTimeMillis();\n";
        String postfix = "\nlong endTime = System.currentTimeMillis();\n";
        String threadName = "\nString threadName = Thread.currentThread().getName();\n";

        String methodName = ctMethod.getName();
        String newMethodName = methodName + "$old";
        ctMethod.setName(newMethodName);

        String returnType = null;
        try {
            returnType = ctMethod.getReturnType().getName();

            if (returnType.trim().toLowerCase().equals("void")) {
                returnType = null;
            }
        } catch (NotFoundException e) {
//            e.printStackTrace();
        }

        CtMethod newMethod = CtNewMethod.copy(ctMethod, methodName, ctClass, null);
        StringBuilder bodyStr = new StringBuilder();
        bodyStr.append("{");
        bodyStr.append(prefix);
        bodyStr.append(threadName);
        // 调用原有代码，类似于method();($$)表示所有的参数
        if (returnType != null) {
            bodyStr.append(returnType + " ret = " + newMethodName + "($$);\n");
        } else {
            bodyStr.append(newMethodName + "($$);\n");
        }
        bodyStr.append(postfix);
        bodyStr.append("System.out.println(\"==============[\" + threadName + \"][" + ctClass.getName() + "][" + methodName + "] Cost:\" +(endTime - startTime) +\"ms " + "===\");\n");
        if (returnType != null) {
            bodyStr.append("return ret;}");
        } else {
            bodyStr.append("}");
        }

        // 替换新方法
        newMethod.setBody(bodyStr.toString());
        // 增加新方法
        ctClass.addMethod(newMethod);
    }
}